"""
@Time: 2021/1/14 下午 4:18
@Author: jinzhuan
@File: el_toolkit.py
@Desc: 
"""

from cognlp.toolkit.base_toolkit import BaseToolkit
import threading
import blink.predictor as predictor
from cognlp.utils.cognet import CognetServer
from cognlp.utils.util import load_json, get_all_forms


class ElToolkit(BaseToolkit):

    def __init__(self, bert_model=None, model_path=None, vocabulary_path=None, device=None,
                 device_ids=None, max_seq_length=None, wikidata2wikipedia_path=None):
        super().__init__(bert_model, model_path, vocabulary_path, device, device_ids, max_seq_length)
        self.wikidata2wikipedia = load_json(wikidata2wikipedia_path)
        self.cognet = CognetServer()
        self.id2url, self.et_ner_model, self.et_models = predictor.get_et_predictor()

    _instance_lock = threading.Lock()

    def __new__(cls, *args, **kwargs):
        if not hasattr(ElToolkit, "_instance"):
            with ElToolkit._instance_lock:
                if not hasattr(ElToolkit, "_instance"):
                    ElToolkit._instance = object.__new__(cls)
        return ElToolkit._instance

    def run(self, sentence):
        url = "https://en.wikipedia.org/wiki/"
        links = predictor.run(10, *self.et_models, text=sentence, id2url=self.id2url, ner_model=self.et_ner_model)
        for link in links:
            forms = get_all_forms(link["title"])
            cognet_link = "unk"
            for form in forms:
                wikipedia = url + form
                if wikipedia in self.wikidata2wikipedia:
                    wikidata = self.wikidata2wikipedia[wikipedia]
                    cognet_link = self.cognet.query("<" + wikidata + ">")
            link["cognet_link"] = cognet_link
        return links

"""
Input:
Ontario is the most populous province in Canada.
Output:
[
    {
        "mention": "ontario",
        "start": 0,
        "end": 1,
        "id": 10915,
        "title": "Ontario",
        "text": " Ontario (; ) is one of the 13 provinces and territories of Canada. Located in Central Canada, it is Canada's most populous province accounting for 38.3 percent of the country's population, and is the second-largest province in total area. Ontario is fourth-largest jurisdiction in total area when the territories of the Northwest Territories and Nunavut are included. It is home to the nation's capital city, Ottawa, and the nation's most populous city, Toronto, which is also Ontario's provincial capital.  Ontario is bordered by the province of Manitoba to the west, Hudson Bay and James Bay to the north, and Quebec to the east and northeast, and to the south by the U.S. states of (from west to east) Minnesota, Michigan, Ohio, Pennsylvania, and New York. Almost all of Ontario's border with the United States follows inland waterways: from the west at Lake of the Woods, eastward along the major rivers and lakes of the Great Lakes/Saint Lawrence River drainage system. These are the Rainy River, the Pigeon River, Lake Superior, the St. Marys River, Lake Huron, the St. Clair River, Lake St. Clair, the Detroit River, Lake Erie, the Niagara River, Lake Ontario and along the St. Lawrence River from Kingston, Ontario, to the Quebec boundary just east of Cornwall, Ontario. There is only about of land border made up of portages including Height of Land Portage on the Minnesota border.  Ontario is sometimes conceptually divided into two regions, Northern Ontario and Southern Ontario. The great majority of Ontario's population and arable land",
        "url": "https://en.wikipedia.org/wiki?curid=22218",
        "cognet_link": "http://cognet.top/entity/E2668966"
    },
    {
        "mention": "canada",
        "start": 7,
        "end": 8,
        "id": 1001810,
        "title": "Canada",
        "text": " Canada ( ) is a country in the northern part of North America. Its ten provinces and three territories extend from the Atlantic to the Pacific and northward into the Arctic Ocean, covering , making it the world's second-largest country by total area. Its southern border with the United States, stretching some , is the world's longest bi-national land border. Canada's capital is Ottawa, and its three largest metropolitan areas are Toronto, Montreal, and Vancouver.  As a whole, Canada is sparsely populated, the majority of its land area being dominated by forest and tundra. Its population is highly urbanized, with over 80 percent of its inhabitants concentrated in large and medium-sized cities, and 70 percent of citizens residing within of the southern border. Canada's climate varies widely across its vast area, ranging from arctic weather in the north, to hot summers in the southern regions, with four distinct seasons.  Various indigenous peoples have inhabited what is now Canada for thousands of years before European colonization. Beginning in the 16th century, British and French expeditions explored and later settled along the Atlantic coast. As a consequence of various armed conflicts, France ceded nearly all of its colonies in North America in 1763. In 1867, with the union of three British North American colonies through Confederation, Canada was formed as a federal dominion of four provinces. This began an accretion of provinces and territories and a process of increasing autonomy from the United Kingdom. This widening autonomy was highlighted by the Statute of Westminster of",
        "url": "https://en.wikipedia.org/wiki?curid=5042916",
        "cognet_link": "http://cognet.top/entity/E2763917"
    }
]
"""